import torch
import torch.nn.functional as F
import cv2
import torchvision.transforms as transforms
from torchvision.utils import save_image

###############################################
## Warning: Do not modify any lines below  ###
###############################################
# Read the image
image = cv2.imread('porche.png')
# Convert BGR image to RGB image
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Define a transform to convert the image to torch tensor
transform = transforms.Compose([transforms.ToTensor()])
# Convert the image to Torch tensor, and make it [B=1, C=3, H, W]
I = transform(image)
I = I.unsqueeze(0)
print("Input image size:", I.size())
###############################################

###############################################
## TODO: move the car to 200 pixels up and left
###############################################

# 平移量：向左上各平移200像素
shift_x = 100  # 水平向左平移量
shift_y = 100  # 垂直向上平移量

# 方法1: 使用2D卷积实现平移
# 创建2D卷积核，在指定位置设置脉冲
H_2d = torch.zeros(1, 1, shift_y+1, shift_x+1).float()
# 在卷积核的右下角设置1，实现向左上平移
H_2d[0, 0, shift_y, shift_x] = 1.0

# 扩展卷积核维度以匹配输入图像的通道数
H_2d_expanded = H_2d.repeat(3, 1, 1, 1)
print(H_2d_expanded.size())

# 使用卷积进行平移（不裁剪，保持原图尺寸）
out1 = F.conv2d(I, H_2d_expanded, padding=0, groups=3)

# 由于平移后图像尺寸变小，我们需要将其填充回原图尺寸
_, _, H, W = I.size()
result_2d = torch.zeros_like(I)
# 将平移后的内容放置在结果图像的左上角
result_h = out1.size(2)
result_w = out1.size(3)
result_2d[:, :, :min(H, result_h), :min(W, result_w)] = out1[:, :, :min(H, result_h), :min(W, result_w)]

save_image(result_2d, 'move_2d.png')


# 方法2: 使用两个1D卷积分别实现平移
# 水平卷积核：向左平移shift_x像素
H_1d_horizon = torch.zeros(1, 1, 1, shift_x+1).float()
H_1d_horizon[0, 0, 0, shift_x] = 1.0  # 在右侧设置1，实现向左平移

# 垂直卷积核：向上平移shift_y像素
H_1d_vertical = torch.zeros(1, 1, shift_y+1, 1).float()
H_1d_vertical[0, 0, shift_y, 0] = 1.0  # 在下方设置1，实现向上平移

# 扩展卷积核维度
H_1d_horizon_expanded = H_1d_horizon.repeat(3, 1, 1, 1)
H_1d_vertical_expanded = H_1d_vertical.repeat(3, 1, 1, 1)

# 先进行水平方向卷积（向左平移）
out2 = F.conv2d(I, H_1d_horizon_expanded, groups=3, stride=1, padding=0)
# 再进行垂直方向卷积（向上